﻿/*
   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 3 of the License, or
   (at your option) any later version.
  
   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.
 
   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software Foundation,
   Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA
*/

using System;
using System.Diagnostics;
using System.Linq;
using System.Runtime.InteropServices;

namespace BHInjector
{
    public class DllInjector
    {
        private readonly Logger logger;

        public DllInjector(Logger logger)
        {
            this.logger = logger;
        }

        public bool Inject(int processId, string dllName)
        {
            if (FindDll(processId, dllName))
            {
                logger.Log("I", "BH already injected, skipping (pid: {0}).", processId);
                return true;
            }

            logger.Log("I", "Injecting...");
            IntPtr hProcess = IntPtr.Zero;
            IntPtr hKernel32 = IntPtr.Zero;
            try
            {
                hProcess = OpenProcess(ProcessAccessFlags.All, false, processId);
                if (hProcess == IntPtr.Zero)
                {
                    logger.Log("I", "Error: cannot open process (error code: {0})", Marshal.GetLastWin32Error());
                    return false;
                }

                hKernel32 = LoadLibrary(@"kernel32.dll");
                if (hKernel32 == IntPtr.Zero)
                {
                    logger.Log("I", "Error: cannot open kernel32.dll (error code: {0})", Marshal.GetLastWin32Error());
                    return false;
                }

                var lpLoadLibraryW = GetProcAddress(hKernel32, @"LoadLibraryW");
                if (lpLoadLibraryW == IntPtr.Zero)
                {
                    logger.Log("I", "Error: cannot get proc address (error code: {0})", Marshal.GetLastWin32Error());
                    return false;
                }

                var lpRemoteString = VirtualAllocEx(hProcess, IntPtr.Zero, dllName.Length + 2, AllocationType.Reserve | AllocationType.Commit, MemoryProtection.ReadWrite);

                if (lpRemoteString != IntPtr.Zero)
                {
                    WriteProcessMemory(hProcess, lpRemoteString, dllName, new IntPtr(dllName.Length * 2 + 2), 0);

                    var hThread = CreateRemoteThread(hProcess, IntPtr.Zero, IntPtr.Zero, lpLoadLibraryW, lpRemoteString, 0, IntPtr.Zero);

                    WaitForSingleObject(hThread, INFINITE);

                    CloseHandle(hThread);
                    VirtualFreeEx(hProcess, lpRemoteString, 0, FreeType.Release);
                }
                else
                {
                    logger.Log("I", "Error: cannot allocate memory (error code: {0})", Marshal.GetLastWin32Error());
                    return false;
                }
            }
            finally
            {
                FreeLibrary(hKernel32);
                CloseHandle(hProcess);
            }

            var dllFound = FindDll(processId, dllName);

            logger.Log("I", dllFound ? "Mission completed! ^_^" : "DLL injection failed! Be sure you are using Diablo 2 1.13c and BH.dll is present.");

            return dllFound;
        }

        private static bool FindDll(int processId, string dllName)
        {
            return Process.GetProcessById(processId).Modules.Cast<ProcessModule>().Any(m => m.ModuleName == dllName);
        }

        #region Helper methods

        [Flags]
        private enum ProcessAccessFlags : uint
        {
            All = 0x001F0FFF,
        }

        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Unicode)]
        static extern IntPtr LoadLibrary(string lpFileName);

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern IntPtr FreeLibrary(IntPtr hModule);

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern IntPtr OpenProcess(ProcessAccessFlags dwDesiredAccess, [MarshalAs(UnmanagedType.Bool)] bool bInheritHandle, int dwProcessId);

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern int CloseHandle(IntPtr hObject);

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern IntPtr GetProcAddress(IntPtr hModule, string lpProcName);

        [Flags]
        public enum AllocationType
        {
            Commit = 0x1000,
            Reserve = 0x2000,
        }

        [Flags]
        public enum MemoryProtection
        {
            ReadWrite = 0x04,
        }

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern IntPtr VirtualAllocEx(IntPtr hProcess, IntPtr lpAddress, int dwSize, AllocationType flAllocationType, MemoryProtection flProtect);

        [Flags]
        public enum FreeType
        {
            Release = 0x8000,
        }

        [DllImport("kernel32.dll", SetLastError = true, ExactSpelling = true)]
        static extern bool VirtualFreeEx(IntPtr hProcess, IntPtr lpAddress, int dwSize, FreeType dwFreeType);

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern int WriteProcessMemory(IntPtr hProcess, IntPtr lpBaseAddress, [MarshalAs(UnmanagedType.LPWStr)]string buffer, IntPtr size, int lpNumberOfBytesWritten);

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern IntPtr CreateRemoteThread(IntPtr hProcess, IntPtr lpThreadAttribute, IntPtr dwStackSize, IntPtr lpStartAddress, IntPtr lpParameter, uint dwCreationFlags, IntPtr lpThreadId);

        const UInt32 INFINITE = 0xFFFFFFFF;

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern UInt32 WaitForSingleObject(IntPtr hHandle, UInt32 dwMilliseconds);

        #endregion
    }
}
